import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import DataLoader, random_split

import torchvision.transforms as transforms
import torchvision.datasets as datasets

import utils
import config
import models.LeNet

def get_loader(batch_size):
    # input of LeNet must be (1x32x32)
    train_transform = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
    test_transform = transforms.Compose([transforms.Resize(32), transforms.ToTensor()])
    full_data = datasets.MNIST(root=config.DatasetDir, train=True, download=True, transform=train_transform)
    test_data = datasets.MNIST(root=config.DatasetDir, train=False, download=True, transform=test_transform)

    # 计算训练集和验证集的大小
    train_size = int(0.8 * len(full_data))
    val_size = len(full_data) - train_size
    # 从完整训练集中随机划分训练集和验证集
    train_data, valid_data = random_split(full_data, [train_size, val_size])

    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(valid_data, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True)
    return train_loader, valid_loader, test_loader

def calc_accuracy(output, target):
    return output.argmax(dim=-1).eq(target).float().mean()

def main():
    utils.set_seed()

    batch_size = config.LeNetConfig.batch_size
    epochs = config.LeNetConfig.epochs
    learning_rate = config.LeNetConfig.learning_rate
    device = config.LeNetConfig.device

    train_loader, valid_loader, test_loader = get_loader(batch_size)

    model = models.LeNet.get_model().to(device)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min',
        factor=config.LeNetConfig.scheduler_factor,
        patience=config.LeNetConfig.scheduler_patience,
    )
    criterion = nn.CrossEntropyLoss()

    # utils.check_net_valid(model, optimizer, criterion, train_loader, epochs, device, calc_accuracy)

    utils.train(model, optimizer, scheduler, criterion, train_loader, valid_loader, epochs, device, calc_accuracy, config.LeNetConfig.valid_patience, config.LeNetConfig.writer)
    utils.test(model, criterion, test_loader, device, calc_accuracy)

    state = {
        "model_dict" : model.state_dict(),
        "optimizer_dict" : optimizer.state_dict(),
        "scheduler_dict" : scheduler.state_dict()
    }
    utils.save_checkpoint(state, config.LeNetConfig.checkpoint_path)

if __name__ == '__main__':
    main()



